from sqlmodel import JSON, Field, SQLModel

from db.constants import StatusEnum

class LLMFactoryBase(SQLModel):
    name: str = Field(primary_key=True, max_length=128, description="LLM factory name")
    logo: str | None = Field(default=None, description="llm logo base64")
    tags: str = Field(max_length=255, index=True, description="LLM, Text Embedding, Image2Text, ASR")
    status: str = Field(default=StatusEnum.VALID.value, max_length=1, index=True, description="is it validate(0: wasted, 1: validate)")

class LLMFactory(LLMFactoryBase, table=True):
    pass

class LLMBase(SQLModel):
    fid: str = Field(primary_key=True, max_length=128, description="LLM factory id", index=True)
    llm_name: str = Field(primary_key=True, max_length=128, description="LLM name", index=True)
    model_type : str = Field(max_length=128, description="LLM, Text Embedding, Image2Text, ASR", index=True)
    max_tokens : int = Field(default=0)
    tags : str = Field(max_length=255, description="LLM, Text Embedding, Image2Text, Chat, 32k...", index=True)
    is_tools : bool = Field(description="support tools", default=False)
    status : str = Field(max_length=1, description="is it validate(0: wasted, 1: validate)", 
                         default=StatusEnum.VALID.value, index=True)

class LLM(LLMBase, table=True):
    pass

class TenantLLMBase(SQLModel):
    name: str = Field(primary_key=True, max_length=128, description="User's LLM name")
    llm_factory: str = Field(max_length=128, description="LLM factory name", index=True)
    llm_name: str = Field(max_length=128, description="LLM name", default="", index=True)
    model_type: str = Field(max_length=128, description="LLM, Text Embedding, Image2Text, ASR", index=True)
    api_key: str | None = Field(default=None, max_length=2048, description="API KEY", index=True)
    api_base: str | None = Field(default=None, max_length=255, description="API Base")
    max_tokens: int = Field(default=8192, index=True)
    used_tokens: int = Field(default=0, index=True)
    
    llm_setting: dict = Field(sa_type=JSON, default={"temperature": 0.1, "top_p": 0.3, "frequency_penalty": 0.7, "presence_penalty": 0.4})

class TenantLLM(TenantLLMBase, table=True):
    pass